import torch.nn.functional as F
import torch.nn as nn

class PositionwiseFeedForward(nn.Module):

    def __init__(self, feature_dim, d_ff, dropout=0.1):
        """d_ff就是第二个线性层的输入维度和第一个线性层的输出维度. """
        super(PositionwiseFeedForward, self).__init__()
        self.w1 = nn.Linear(feature_dim, d_ff)
        self.w2 = nn.Linear(d_ff, feature_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w2(self.dropout(F.relu(self.w1(x))))